# 展示 color  coord  segment.npy 的展示效果

# 展示点云分割 显示效果

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import os
import open3d as o3d


def generate_label_colors(labels, color_map):
    """为标签生成对应的颜色"""
    colors = np.zeros((len(labels), 3))
    for label, color in color_map.items():
        colors[labels == label] = color
    return colors


def visualize_three_views(pcd_original, pcd_gt, title1, title2,):
    """同时可视化原始点云、GT和预测结果, 并同步视角
    
    参数:
        pcd_original (o3d.geometry.PointCloud): 原始点云 (RGB颜色)
        pcd_gt (o3d.geometry.PointCloud): Ground Truth点云(伪彩色)
        title1 (str): 第一个窗口标题（原始点云）
        title2 (str): 第二个窗口标题 (GT点云))
    """
    # 创建三个窗口并设置标题
    vis1 = o3d.visualization.Visualizer()
    vis1.create_window(window_name=title1, width=800, height=600, left=0)  # 左侧窗口
    vis1.add_geometry(pcd_original)
    
    vis2 = o3d.visualization.Visualizer()
    vis2.create_window(window_name=title2, width=800, height=600, left=800)  # 中间窗口
    vis2.add_geometry(pcd_gt)
    

    # 同步视角（以第一个窗口为基准）
    params = vis1.get_view_control().convert_to_pinhole_camera_parameters()
    vis2.get_view_control().convert_from_pinhole_camera_parameters(params)
    
    # 运行可视化
    print("按 Q 或关闭窗口退出...")
    while True:
        if not (vis1.poll_events() and vis2.poll_events()):
            break
        vis1.update_renderer()
        vis2.update_renderer()
    
    # 关闭窗口
    vis1.destroy_window()
    vis2.destroy_window()


def visualize_segmentation(data_dir):
    # 定义标签名称和对应的值
    label_info = [
        ("Eyebrows ", 0),
        ("Eyes", 1),
        ("Nose", 2),
        ("Mouth", 3),
        ("Ears", 4),
        ("Hair", 5),
        ("Neck", 6),
        ("Bracket", 7),
        ("Face", 8),
        ("Background", 9),
        ("ignore", -1)
    ]
    
    names = [item[0] for item in label_info]
    label_values = [item[1] for item in label_info]
    
    # 为每个标签分配不同的颜色
    colors = [
        [151, 223, 137],    # wall
        [174, 200, 232],    # floor
        [31, 120, 180],     # cabinet
        [255, 189, 120],    # bed
        [189, 190, 34],     # chair
        [140, 86, 74],      # sofa
        [255, 152, 151],    # table
        [213, 39, 40],      # door
        [196, 177, 213],    # window
        [148, 103, 189],    # bookshelf
        [23, 24, 22],       # ignore (黑色)
    ]
    
    # 创建标签到颜色的映射
    color_map = {value: np.array(color)/255.0 for value, color in zip(label_values, colors)}
     
    # 加载数据
    # scene_dir = os.path.join(data_dir, 'scene0050_00')
    color_data = np.load(os.path.join(data_dir, 'output_color.npy'))
    coord_data = np.load(os.path.join(data_dir, 'output_coord.npy'))
    gt_data = np.load(os.path.join(data_dir, 'output_segment.npy'))

    # 2. 3D点云可视化部分
    # 合并坐标和颜色
    original_point_cloud = np.hstack([coord_data, color_data])
    print(original_point_cloud)
    
    # 可视化原始点云和预测标签点云
    visualize_point_cloud_with_labels(
        original_point_cloud, 
        gt_data,
        color_map,
        "原始点云 (RGB)", 
        "Ground Truth",
    )


def visualize_point_cloud_with_labels(points_with_color, gt_label, color_map, title1, title2):
    """可视化点云和标签"""
    # 检查数据格式和一致性
    # assert len(points_with_color) == len(pred_labels), "点云和标签数量不匹配!"
    # print("原始点云形状:", points_with_color.shape)
    # print("标签形状:", pred_labels.shape)

    # 提取坐标和原始颜色
    points = points_with_color[:, :3]
    original_colors = points_with_color[:, 3:6] / 255.0

    # # 为预测标签生成伪彩色
    label_colors = generate_label_colors(gt_label, color_map)

    # 创建Open3D点云对象
    pcd_original = o3d.geometry.PointCloud()
    pcd_original.points = o3d.utility.Vector3dVector(points)
    pcd_original.colors = o3d.utility.Vector3dVector(original_colors)

    pcd_gt = o3d.geometry.PointCloud()
    pcd_gt.points = o3d.utility.Vector3dVector(points)
    pcd_gt.colors = o3d.utility.Vector3dVector(generate_label_colors(gt_label, color_map))

    pcd_label = o3d.geometry.PointCloud()
    pcd_label.points = o3d.utility.Vector3dVector(points)
    pcd_label.colors = o3d.utility.Vector3dVector(label_colors)

    # 可选：下采样（如果点云过大）
    if len(points) > 100000:
        voxel_size = 0.03
        pcd_original = pcd_original.voxel_down_sample(voxel_size)
        pcd_label = pcd_label.voxel_down_sample(voxel_size)

    # 可视化
    visualize_three_views(pcd_original, pcd_gt, title1, title2)


# 使用示例
data_directory = 'D:\\resource\\code'  # 替换为你的数据文件夹路径
visualize_segmentation(data_directory)
